# import os
# import cv2
# import sys
# import time
# import numpy as np
# import onnxruntime as ort  # 用于加载和推理 ONNX 模型
# from label_zh import coco_zh  # 确保该文件存在，包含 COCO 类别中文映射

# def parse_onnx_output(outputs):
#     """解析 ONNX 模型输出，逻辑与 RKNN 后处理保持一致"""
#     print("输出张量形状:", outputs[0].shape)  # YOLOv8 输出通常为 (1, 84, 8400)
#     predictions = np.squeeze(outputs[0])  # 去除 batch 维度 -> (84, 8400)
#     print("去除batch后形状:", predictions.shape)

#     # 若形状为 (8400, 84)，转置为 (84, 8400)（适配不同导出参数）
#     if predictions.shape == (8400, 84):
#         predictions = predictions.T

#     # 提取框坐标、置信度和类别 ID
#     boxes = predictions[:4, :]  # 前4行：x_center, y_center, width, height（相对坐标）
#     scores = predictions[4:, :].max(axis=0)  # 类别置信度（取最大值）
#     print("最大置信度:", np.max(scores))
#     print("置信度大于0.5的数量:", np.sum(scores > 0.5))

#     class_ids = predictions[4:, :].argmax(axis=0)  # 最高置信度对应的类别 ID

#     # 过滤低置信度结果（阈值 0.1）
#     mask = scores > 0.5
#     boxes = boxes[:, mask]
#     scores = scores[mask]
#     class_ids = class_ids[mask]

#     # 转换为结构化标签（与 RKNN 输出格式一致）
#     labels = []
#     for i in range(len(scores)):
#         labels.append({
#             "label": coco_zh[int(class_ids[i])],
#             "category": "物体",
#             "confidence": float(scores[i]),
#             "description": f"检测到{coco_zh[int(class_ids[i])]}"
#         })

#     return labels

# def simulate_onnx_recognition(image_path, onnx_model_path):

#     # 1. 加载 ONNX 模型耗时
#     load_start = time.time()
#     # 初始化 ONNX Runtime 会话（使用 CPU 推理，可指定 providers 为 GPU 加速）
#     session = ort.InferenceSession(
#         onnx_model_path,
#         providers=['CPUExecutionProvider']  # 若有 GPU 可添加 ['CUDAExecutionProvider']
#     )
#     # 获取输入节点名称（YOLOv8 通常为 'images'）
#     input_name = session.get_inputs()[0].name
#     input_shape = session.get_inputs()[0].shape  # 通常为 (1, 3, 640, 640)
#     print(f"ONNX 模型输入形状: {input_shape}")
#     load_end = time.time()
#     print(f"模型加载耗时: {(load_end - load_start)*1000:.2f} ms")

#     # 2. 图像预处理耗时（与 RKNN 预处理保持一致）
#     img = cv2.imread(image_path)
#     if img is None:
#         raise FileNotFoundError(f"未找到图像文件: {image_path}")
#     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # BGR -> RGB
#     img = cv2.resize(img, (640, 640))  # 缩放至模型输入尺寸
#     img = img / 255.0  # 归一化到 0-1 范围
#     # ONNX 模型通常期望 NCHW 格式（与 RKNN 的 NHWC 不同，需注意）
#     img = img.transpose(2, 0, 1)  # 转换为 (3, 640, 640)（CHW）
#     img = np.expand_dims(img, axis=0).astype(np.float32)  # 最终形状 (1, 3, 640, 640)

#     # 保存预处理后的图像用于验证
#     preprocessed_img = (img[0].transpose(1, 2, 0) * 255).astype(np.uint8)  # 转回 HWC
#     preprocessed_img = cv2.cvtColor(preprocessed_img, cv2.COLOR_RGB2BGR)
#     cv2.imwrite("onnx_preprocessed_img.jpg", preprocessed_img)

#     # 输入数据需与模型期望格式一致（NCHW）
#     outputs = session.run(None, {input_name: img})
#     print(f"推理输出形状:", [o.shape for o in outputs])
#     print(f"推理输出最大值:", [np.max(o) for o in outputs])
#     print(f"推理输出最小值:", [np.min(o) for o in outputs])
#     # 4. 后处理
#     labels = parse_onnx_output(outputs)

#     return labels

# if __name__ == '__main__':
#     # 测试图像路径和 ONNX 模型路径（根据实际情况修改）

#     pwd_dir = os.path.dirname(os.path.abspath(__file__))
#     image_path = os.path.join(pwd_dir, './bus.jpg')
#     onnx_model_path = os.path.join(pwd_dir, './yolov8n.onnx')
# #    onnx_model_path = 'yolov8n.onnx'  # 替换为你的 YOLOv8n ONNX 模型路径

#     # 执行 ONNX 模型推理
#     label = simulate_onnx_recognition(image_path, onnx_model_path)
#     # print("my_label:")
#     # print(label)



from ultralytics import YOLO
import os
import cv2  # 用于图像读取、绘制和保存
import numpy as np  # 辅助图像处理

# 1. 定义COCO数据集类别中文名称（补全之前的英文，确保全中文）
coco_zh = [
    "persson", "自行车", "汽车", "摩托车", "飞机", "公交车", "火车", "卡车", "船",
    "交通灯", "消防栓", "停车标志", "长椅", "鸟", "猫", "狗", "马", "羊", "牛",
    "大象", "熊", "斑马", "长颈鹿", "背包", "雨伞", "手提包", "领带", "行李箱",  # 修正：suitcase→行李箱
    "飞盘", "滑雪板", "滑雪杖", "运动球", "风筝", "棒球棒", "棒球手套", "滑板",
    "冲浪板", "网球拍", "瓶子", "红酒杯", "杯子", "叉子", "刀", "勺子", "碗",
    "香蕉", "苹果", "三明治", "橙子", "西兰花", "胡萝卜", "热狗", "披萨",  # 修正：broccoli→西兰花
    "甜甜圈", "蛋糕", "椅子", "沙发", "盆栽", "床", "餐桌", "厕所", "电视",
    "笔记本电脑", "鼠标", "遥控器", "键盘", "手机", "微波炉", "烤箱", "烤面包机",
    "水槽", "冰箱", "书", "时钟", "花瓶", "剪刀", "泰迪熊", "吹风机", "牙刷"
]


if __name__ == '__main__':
    print("start test.py")
    pwd_dir = os.path.dirname(os.path.abspath(__file__))
    image_path = os.path.join(pwd_dir, 'bus.jpg')  # 原始图片路径
    model_path = os.path.join(pwd_dir, 'yolov8n.pt')  # 模型路径
    output_image_path = os.path.join(pwd_dir, 'detected_bus.jpg')  # 生成的标注后图片路径
    print(f"原始图片路径: {image_path}")
    print(f"模型路径: {model_path}")
    print(f"标注后图片保存路径: {output_image_path}")

    # 2. 加载模型并推理
    model = YOLO(model_path)
    results = model.predict(
        image_path,
        save=False,      # 不使用YOLO默认保存，自定义保存逻辑
        save_txt=False,
        save_conf=False,
        save_crop=False,
        exist_ok=False,
        imgsz=640  # 推理输入尺寸（可选，默认640）
    )
    print("YOLO end")
    print(results)
    # 3. 读取原始图片（OpenCV默认BGR格式，后续绘制用）
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片，请检查路径: {image_path}")

    # 4. 遍历检测结果，画框+写标签
    labels = []  # 用于存储结构化标签
    # 定义画框颜色（BGR格式，这里用绿色，可自定义）
    box_color = (0, 255, 0)  # 绿色：(蓝, 绿, 红)
    # 定义标签文字颜色（白色）和字体
    text_color = (255, 255, 255)
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.5  # 字体大小
    line_thickness = 2  # 框线厚度

    for result in results:
        for box in result.boxes:
            # 提取检测框坐标（YOLO输出的是归一化坐标，需转换为像素坐标）
            # box.xyxy: [x1, y1, x2, y2]（左上角x/y，右下角x/y，归一化后的值）
            x1, y1, x2, y2 = box.xyxy[0].tolist()  # 转换为列表（像素坐标）
            # 转换为整数（像素坐标必须是整数）
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            
            # 提取类别ID和置信度
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            cls_name = coco_zh[cls_id]  # 获取中文类别名

            # 4.1 绘制检测框（在图片上画矩形）
            cv2.rectangle(img, (x1, y1), (x2, y2), box_color, line_thickness)

            # 4.2 绘制标签文字（格式：类别名: 置信度）
            text = f"{cls_name}: {conf:.2f}"  # 保留2位小数的置信度
            # 计算文字位置（在框的左上角，避免超出图片边界）
            text_x = x1
            text_y = y1 - 10 if y1 > 10 else y1 + 20  # 防止文字超出图片顶部
            # 先画文字背景（黑色半透明，提升可读性）
            text_size = cv2.getTextSize(text, font, font_scale, line_thickness)[0]
            text_bg_x1 = text_x
            text_bg_y1 = text_y - text_size[1]
            text_bg_x2 = text_x + text_size[0]
            text_bg_y2 = text_y
            cv2.rectangle(img, (text_bg_x1, text_bg_y1), (text_bg_x2, text_bg_y2), (0, 0, 0), -1)  # 黑色背景
            # 再写文字
            cv2.putText(img, text, (text_x, text_y), font, font_scale, text_color, line_thickness)

            # 4.3 存储结构化标签（与之前逻辑一致）
            labels.append({
                "label": cls_name,
                "category": "物体",
                "confidence": conf,
                "bbox": [x1, y1, x2, y2],  # 像素坐标的检测框
                "description": f"检测到{cls_name}"
            })

    # 5. 保存标注后的图片
    cv2.imwrite(output_image_path, img)
    print(f"标注后图片已保存至: {output_image_path}")

    # 6. 打印结构化标签
    print("\n检测结果标签:")
    for idx, label in enumerate(labels, 1):
        print(f"{idx}. {label['description']} | 置信度: {label['confidence']:.4f} | 坐标: {label['bbox']}")
